import ImageMetrics as IM
import pylab
import Image
from numpy import zeros, mgrid

import JT
import Calibration

#===============================================================================
#
#    Data Paths
#
#===============================================================================

fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
#fixedImageName = 'C:/Users/bryan/bryan-code/trunk/Images/CalibratedDRRImage.mhd'
fixedImageMaskName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ucharb2Rmask080-LAT.mha'

inputVolFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mha'
maskFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1b2VOI.mha'
staFile = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
calFileLat = 'C:/Users/bryan/bryan-code/trunk/TestData/ext_cal1.txt'


#===============================================================================
#
#    Setup DRR
#
#===============================================================================

xraycam = JT.XRayCamera()
drr = JT.DRR()
drr.SetXRayCamera(xraycam)
drr.SetVolumeFilename(inputVolFileName)
drr.SetBlackOnWhiteImage()
drr.InteractionOff()

cal = Calibration.ExternalCalibration()
cal.LoadConsolidatedCalibration(calFileLat)
cal.LoadStaFile(staFile)

xraycam.SetExternalCalibration(cal)

drr._volume.SetOriginalTransform(cal._VolumeT)
drr._volume.SetBodyCenter((3.906, 2.17, -7.378))

# Set Fixed image using path
fixedImage = JT.FixedImage()
fixedImage.SetFileName(fixedImageName)

reg=JT.Registration()
reg.SetFixedImage(fixedImage.GetImage(0))
reg.SetMovingImageGenerator(drr)
reg._blur = 2.0

#Set Mask (choose one, moving or fixed)
#drr._volume.SetImageMask(maskFileName)
reg.SetFixedImageMask(fixedImageMaskName)

# Setup grid to evaluate over
xdim = 2
ydim = 2
m = mgrid[-xdim:xdim+1,-ydim:ydim+1]

gridx = m[0,:,:].flatten()
gridy = m[1,:,:].flatten()
scale = 1

val = zeros((2*xdim+1,2*ydim+1))
iter = 0

if 1:
    for i in range(-xdim,xdim+1):
        for j in range(-ydim,ydim+1):
#            pylab.imshow(currentImage)
            pose = [i*scale,j*scale,0,0,0,0]
            val[i][j] = reg.GetValue(pose)
            print "Iteration: %i [%i,%i], Metric Value: %f" % (iter,i*scale,j*scale,val[i][j])
            iter += 1
        
if 0:
    for i in range(-gridx/2,gridx/2):
        movingImage=fixedImage.rotate(i*2)
        val[i][j] = IM.GD(fixedImage, movingImage)
        print "Iteration: %i, Metric Value: %f" % (iter,val[i][j])
        iter += 1

pylab.figure

pylab.imshow(val)
